﻿#pragma once

#include "coroutine/coroutine.h"
#include "rpc_defines.h"
#include "rpc_logger.h"
#include "rpc_service.h"
#include "rpc_singleton.h"
#include "rpc_statistics.hh"
#include "rpc_stub.h"
#include "rpc_transport.h"
#include <ctime>
#include <functional>
#include <istream>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <unordered_map>

namespace kratos {
namespace service {
class ServiceContext;
}
} // namespace kratos

namespace rpc {

class Rpc;
class RpcImpl;
class RpcProbe;
class ProxyCreator;
class ProxyManager;
class StubManager;
class ProxyCallManager;
class StubCallManager;
class InterfaceDescriptorFactory;

/**
 * memory allocation
 */
using RpcMalloc = std::function<void *(std::size_t)>;
/**
 * memory deallocation
 */
using RpcFree = std::function<void(void *)>;

// RPC framework
class Rpc {
  /**
   * std::shared_ptr deleter
   */
  template <typename T> void sharedPtrDeleter(T *obj) {
    if (!obj) {
      return;
    }
    auto *typeObject = dynamic_cast<T *>(obj);
    typeObject->~T();
    free(typeObject);
  }

public:
  virtual ~Rpc(){};
  // Initialize
  // @param logger external @see Logger interface pointer
  // @param coroMode Start coroutine for each service implementation
  // @param openStat true means open statistic
  virtual void initialize(Logger *logger = nullptr, bool coroMode = true,
                          bool openStat = false) = 0;
  // Deinitialize
  virtual void deinitialize() = 0;
  // Setup allocator function
  virtual void setAllocator(RpcMalloc mallocFunc, RpcFree freeFunc) = 0;
  // Call this method in your main loop
  // @param ms current timestamp in millionseconds
  virtual void update(std::time_t ms = 0) = 0;
  // Call it when received some RPC data
  // @param transport User implemented Transport interface
  // @param ctx ServiceContext instance
  // @param proxyHandler ProxyHandle for proxy mode
  // @param nonRpcHandler The handler for non-RPC message
  virtual void onMessage(TransportPtr &transport,
                         kratos::service::ServiceContext *ctx = nullptr,
                         rpc::ProxyHandler *proxyHandler = nullptr,
                         rpc::NoRpcHander *nonRpcHandler = nullptr) = 0;
  // Returns @see Logger interface pointer
  virtual Logger *getLogger() = 0;
  // Set user context pointer
  virtual void setContext(void *context) = 0;
  // Get user context pointer
  template <typename T> T *getContext() {
    return reinterpret_cast<T *>(getContextRaw());
  }
  // Relay service call to internal grid
  // @param from The message source
  // @param to The message destination
  // @param header RPC protocol header
  // @param replace_call_id ProxyHandler could make its own call ID to replace
  // real call ID
  // @return true或false
  virtual bool relay(TransportPtr &from, TransportPtr &to,
                     const RpcMsgHeader &header,
                     CallID replace_call_id = INVALID_CALL_ID) = 0;
  /**
   * Returns ProxyManager
   */
  virtual ProxyManager *getProxyManager() = 0;
  /**
   * Returns StubManager
   */
  virtual StubManager *getStubManager() = 0;
  /**
   * Returns ProxyCallManager
   */
  virtual ProxyCallManager *getProxyCallManager() = 0;
  /**
   * Returns StubCallManager
   */
  virtual StubCallManager *getStubCallManager() = 0;
  /**
   * New a std::shared_ptr with customized deleter
   */
  template <typename ObjectType, typename... Args>
  std::shared_ptr<ObjectType> make_shared_ptr(Args... args) {
    auto *memory = malloc(sizeof(ObjectType));
    auto objectPtr = new (memory) ObjectType(std::forward<Args>(args)...);
    return std::shared_ptr<ObjectType>(
        objectPtr, std::bind(&Rpc::sharedPtrDeleter<ObjectType>, this,
                             std::placeholders::_1));
  }
  /**
   * New a object from memory pool
   */
  template <typename ObjectType, typename... Args>
  ObjectType *allocObject(Args... args) {
    auto *memory = malloc(sizeof(ObjectType));
    return new (memory) ObjectType(std::forward<Args>(args)...);
  }
  /**
   * Recycle a object to memory pool
   */
  template <typename ObjectType> void deallocObject(ObjectType *&object) {
    if (!object) {
      return;
    }
    object->~ObjectType();
    free(object);
    object = nullptr;
  }
  /**
   * Allocate a memory block
   * \param size block size in bytes
   */
  virtual void *malloc(std::size_t size) = 0;
  /**
   * Recycle a object to memory pool
   */
  virtual void free(void *p) = 0;
  /**
   * Returns global proxy creator
   */
  virtual ProxyCreator *getProxyCreator() = 0;
  /**
   * Switch On/Off of RPC statistics
   */
  virtual bool switchRpcStat(bool on) = 0;
  /**
   * @brief Setup InterfaceDescriptorFactory
   * @param descriptorFactory InterfaceDescriptorFactory pointer
   */
  virtual void
  setDescriptorFactory(InterfaceDescriptorFactory *descriptorFactory) = 0;
  /**
   * @brief get InterfaceDescriptorFactory
   * @return InterfaceDescriptorFactory pointer
   */
  virtual InterfaceDescriptorFactory *getDescriptorFactory() = 0;
  /**
   * @brief Setup probe
   * @param probe RpcProbe implementation
   */
  virtual void setRpcProbe(RpcProbe *probe) = 0;
  /**
   * @brief 获取RpcProbe
   * @return RPC probe instance
   */
  virtual RpcProbe *getRpcProbe() = 0;

private:
  /**
   * Returns internal context
   */
  virtual void *getContextRaw() = 0;
};

// RPC framework
class RpcImpl : public Rpc {
  bool coroMode_{true};    ///< Start coroutine for each service implementation
  void *context_{nullptr}; ///< User context pointer
  bool openStat_{false};   ///< Is open stat or not

  ProxyManager *proxyManager_{nullptr};         ///< ProxyManager instance
  StubManager *stubManager_{nullptr};           ///< StubManager instance
  ProxyCallManager *proxyCallManager_{nullptr}; ///< ProxyCallManager instance
  StubCallManager *stubCallManager_{nullptr};   ///< StubCallManager instance

  RpcMalloc rpcMalloc_{nullptr}; ///< Memory allocation function
  RpcFree rpcFree_{nullptr};     ///< Memory deallocation function
  RpcProbe *rpcProbe_{nullptr};  ///< RPC probe
  InterfaceDescriptorFactory *descriptorFactory_{
      nullptr}; ///< InterfaceDescriptorFactory pointer

public:
  // ctor
  RpcImpl();
  // dtor
  virtual ~RpcImpl();
  // Initialize
  // @param logger external @see Logger interface pointer
  // @param coroMode Start coroutine for each service implementation
  // @param openStat true means open statistic
  virtual void initialize(Logger *logger = nullptr, bool coroMode = true,
                          bool openStat = false) override;
  // Deinitialize
  virtual void deinitialize() override;
  // Setup allocator function
  virtual void setAllocator(RpcMalloc mallocFunc, RpcFree freeFunc) override;
  // Call this method in your main loop
  // @param ms current timestamp in millionseconds
  virtual void update(std::time_t ms = 0) override;
  // Call it when received some RPC data
  // @param transport User implemented Transport interface
  virtual void onMessage(TransportPtr &transport,
                         kratos::service::ServiceContext *ctx = nullptr,
                         rpc::ProxyHandler *proxyHandler = nullptr,
                         rpc::NoRpcHander *nonRpcHandler = nullptr) override;
  // Returns @see Logger interface pointer
  virtual Logger *getLogger() override;
  // Set user context pointer
  virtual void setContext(void *context) override;
  // Relay service call to internal grid
  virtual bool relay(TransportPtr &from, TransportPtr &to,
                     const RpcMsgHeader &header,
                     CallID replace_call_id = INVALID_CALL_ID) override;

  virtual ProxyManager *getProxyManager() override;
  virtual StubManager *getStubManager() override;
  virtual ProxyCallManager *getProxyCallManager() override;
  virtual StubCallManager *getStubCallManager() override;

  virtual void *malloc(std::size_t size) override;
  virtual void free(void *p) override;
  virtual ProxyCreator *getProxyCreator() override;
  virtual bool switchRpcStat(bool on) override;
  virtual void
  setDescriptorFactory(InterfaceDescriptorFactory *descriptorFactory) override;
  virtual InterfaceDescriptorFactory *getDescriptorFactory() override;
  virtual void setRpcProbe(RpcProbe *probe) override;
  virtual RpcProbe *getRpcProbe() override;

private:
  // Get user context pointer
  virtual void *getContextRaw() override;
  // Call a stub service
  // @param header message header
  // @param transport message stream
  void onCall(const RpcMsgHeader &header, TransportPtr &transport,
              kratos::service::ServiceContext *ctx = nullptr);
  // Process a proxy call
  // @param header message header
  // @param transport message stream
  void onProxyCall(const RpcMsgHeader &header, TransportPtr &transport,
                   kratos::service::ServiceContext *ctx = nullptr);
  // Call a proxy to receive returning message
  // @param header message header
  // @param transport message stream
  void onReturn(const RpcMsgHeader &header, TransportPtr &transport);
  // Process a proxy cakk returning message
  // @param header message header
  // @param transport message stream
  void onProxyCallReturn(const RpcMsgHeader &header, TransportPtr &transport);
  // Build header from stream
  // @param transport message stream
  // @param [OUT] header message header
  // @retval true success
  // @retval false failure
  bool readHeader(TransportPtr &transport, RpcMsgHeader &header);
  // Service not found and return message
  // @param transport message stream
  // @param callID proxy call ID
  void notFound(TransportPtr &transport, CallID callID);
  // Process message
  // @param header message header
  // @param transport message stream
  // @param ctx ServiceContext
  // @param proxyHandler ProxyHandler
  // @param nonRpcHandler The handler for non-RPC message
  void doMessage(const RpcMsgHeader &header, TransportPtr &transport,
                 kratos::service::ServiceContext *ctx = nullptr,
                 rpc::ProxyHandler *proxyHandler = nullptr,
                 rpc::NoRpcHander *nonRpcHandler = nullptr);
  // Find service
  // @param serviceID service instance ID
  // @param uuid service UUID
  StubPtr getService(ServiceID serviceID, ServiceUUID uuid,
                     kratos::service::ServiceContext *ctx = nullptr);
  // Is stub running in a coroutine
  // @retval true yes
  // @retval false no
  bool isCoroMode() { return coroMode_; }
  // Is open statistics or not
  bool isOpenStat() { return openStat_; }

  friend class StubImpl;
  friend class StubCallImpl;
  friend class TransportGuard;
  friend class RelayBuffer;
};

// Destory RPC framework
void deinitialize();

// Load a bundle
// @param uuid Service UUID
// @param bundlePath bundle path
// @param noCreate donot create instance
// @param ctx context
// @retval true success
// @retval false failure
bool loadClass(Rpc *rpc, const std::string &uuid, const std::string &bundlePath,
               bool noCreate = false,
               kratos::service::ServiceContext *ctx = nullptr);
// Destroy all service with service UUID is uuid
bool destroyServiceUnsafe(Rpc *rpc, const std::string &uuid);
// Unload all bundles
bool unloadClassUnsafe(Rpc *rpc, const std::string &uuid);
// Make RPC return
template <typename T, typename... Args>
std::shared_ptr<T> make_return(rpc::StubCallPtr call, Args... args) {
  return call->getRpc()->make_shared_ptr<T>(std::forward<Args>(args)...);
}
// Get service from peer
// @param stubCall RPC call from peer
template <typename T> std::shared_ptr<T> getService(rpc::StubCallPtr stubCall) {
  return rpc::getService<T>(stubCall->getGlobalIndex(),
                            stubCall->getTransport(), true);
}

} // namespace rpc
